#!/usr/bin/python

__author__    = "Anoop Menon <codelogic at gmail dot com>"
__copyright__ = "Copyright (c) 2007"
__license__   = "GPL Version 2.0"
__version__   = "0.8"


#########################################################################
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# If you have any comments or would like to suggest improvements, please
# email me at 'codelogic at gmail dot com'
#
#########################################################################


#########################################################################
#
# Changelog
#
# 11/24/07 - fixed bug in XML filename (thx to h9z3zzj02)
#
# 11/20/07 - fixed bug in path seperator (thx to Kamil Sarniak for
#            pointing it out).
#
# 11/19/07 - added support to recreate an XML which can be fed to the
#            theme compiler to recreate original p3t file.
#
# 11/19/07 - multiple bgimages are correctly extracted instead of being
#            overwritten.
#
# 11/19/07 - removed requirement of having p3tcompiler.exe in the same
#            directory.
#
# 11/16/07 - added support to save GIM files as PNGs
#
# 11/16/07 - added a basic GIM file reader
#
# 11/16/07 - non 'bgimages' are zlib compressed. Added support for that
#            so that icons are decompressed to GIM files.
#
# 11/15/07 - First release
#
#########################################################################

import zipimport
import sys
import struct
import os
import os.path
import traceback
import zlib
import xml.dom.minidom

g_strings = []
g_PIL = 0

try:
    from PIL import Image
    g_PIL = 1
except:
    pass

class GimFile:
    '''
    GIM files (created by the theme compiler) have a 128 byte header followed by
    raw RGBA data width and height are stored at offsets 72 and 74 as short ints
    (big endian)
    '''
    def __init__(self, imagedata):
        self.buf = imagedata
        if self.buf[1:4]!="GIM":
            raise Exception("This data is not in GIM format")
        self.header = self.buf[0:128]
        (self.width, self.height) = struct.unpack(">hh", self.header[72:76])
        self.buf = self.buf[128:]

    def save(self, dest):
        global g_PIL
        if g_PIL==0:
            raise Exception("Python Imaging Library is required to save GIM files in other formats")
        try:
            img = Image.frombuffer("RGBA", (self.width, self.height), self.buf, 'raw', 'RGBA', 0, 1)
            img.save(dest)
        except Exception, e:
            raise e


class P3THeader:
    def __init(self):
        self.magic = ""
        self.version = 0

        self.tree_offset = 0
        self.tree_size = 0

        self.idtable_offset = 0
        self.idtable_size = 0

        self.stringtable_offset = 0
        self.stringtable_size = 0

        self.intarray_offset = 0
        self.intarray_size = 0

        self.floatarray_offset = 0
        self.floatarray_size = 0

        self.filetable_offset = 0
        self.filetable_size = 0


    def __str__(self):
        return '''
        magic:               "%s"
        version:             %d
        tree_offset:         %d
        tree_size:           %d
        idtable_offset:      %d
        idtable_size:        %d
        stringtable_offset:  %d
        stringtable_size:    %d
        intarray_offset:     %d
        intarray_size:       %d
        floatarray_offset:   %d
        floatarray_size:     %d
        filetable_offset:    %d
        filetable_size:      %d

        ''' % (self.magic, self.version, self.tree_offset,
               self.tree_size, self.idtable_offset, self.idtable_size, self.stringtable_offset,
               self.stringtable_size, self.intarray_offset, self.intarray_size,
               self.floatarray_offset, self.floatarray_size, self.filetable_offset,
               self.filetable_size)


class P3TElement(xml.dom.minidom.Element):
    def __init__(self, header):
        xml.dom.minidom.Element.__init__(self, "")
        self.numattr  = 0
        self.stringoffset = 0
        self.parentoffset = 0
        self.name = ""
        self.offset = 0
        self.t = []
        self.attributes = {}
        self.has_file = 0
        self.filename = ""
        self.header = header
        self.elemap = {}

    def __str__(self):
        return ("Name: '%s', no. of attributes: %d" % (self.name, self.numattr)) #+ "\n"+str(self.t)

    def add_subelement(self, ele):
        self.elemap[ele.offset] = ele
        self.appendChild(ele)
        
    def add_attribute(self, attr):
        if attr.type==6:
            self.has_file = 1
        self.attributes[attr.name] = attr
        if attr.name=="size":
            return
        if attr.name and attr.value:
            self.setAttribute(str(attr.name), str(attr.value))

    def dump_files(self, rootdir, hfile):
        for subele in self.elemap.keys():
            self.elemap[subele].dump_files(rootdir, hfile)

        if not self.has_file:
            return
        try:
            os.mkdir(rootdir)
        except:
            pass
        if not os.path.exists(rootdir):
            raise Exception("Unable to create directory: %s" % rootdir)

        fileroot = rootdir + os.path.sep

        pos = hfile.tell()
        fileoffset = self.header.filetable_offset
        
        for k,v in self.attributes.iteritems():
            if v.type==6:
                ext = ".gim"
                suffix = 1
                if self.name=="bgimage":
                    ext = ".jpg"
                while True:
                    try:
                        filename = fileroot + self.attributes['id'].value + "_" + str(suffix) + ext
                    except:
                        filename = fileroot + v.name + "_" + str(suffix) + ext
                    if os.path.exists(filename):
                        suffix+=1
                    else:
                        break
                hfile.seek(fileoffset + v.fileoffset)
                outfile = open(filename, "wb")
                print "Writing original : '%s'" % filename
                cbuf = hfile.read(v.filesize)
                self.filename = filename
                if ext==".gim":
                    # zlib compressed image
                    dbuf = zlib.decompress(cbuf)
                    try:
                        gimfile = GimFile(dbuf)
                        print "Coverting to PNG : '%s' : " % (filename+".png"),
                        gimfile.save(filename+".png")
                        self.filename = filename+".png"
                        print "OK"
                    except Exception, e:
                        print "Failed"
                        print str(e)
                    outfile.write(dbuf)
                else:
                    outfile.write(cbuf)
                outfile.close()

                # remove
                self.setAttribute(str(v.name), os.path.basename(str(self.filename)))                                

        hfile.seek(pos)

    def parse(self, data, format, hfile):
        global g_strings

        f = hfile

        self.offset = hfile.tell() - struct.calcsize(format) - self.header.tree_offset

        t = struct.unpack(format, data)
        self.t = t
        # 1st element is offset into string table
        self.stringoffset = t[0]

        # 2nd element is the no. of attributes
        self.numattr = t[1]

        # 3rd element is the offset of the parent element
        self.parentoffset = t[2]

        pos = f.tell()
        f.seek(self.header.stringtable_offset+self.stringoffset)
        self.name = ""
        a = f.read(1)
        while a!='\x00':
            self.name = self.name + str(a)
            a = f.read(1)
        f.seek(pos)
        self.tagName = self.name

class P3TAttribute:
    def __init__(self):
        self.type  =0
        self.handle = 0
        self.offset = 0
        self.size = 0
        self.value = 0
        self.name = ""
        self.fileoffset = 0
        self.filesize = 0
        self.id = 0
        self.t = []

    def __str__(self):
        #return "type: %d, name: '%s', value: '%s', offset: %d, handle: %d, size: %d" % (self.type, str(self.name), str(self.value), self.offset, self.handle, self.size)
        return ("type: %d, name: '%s', value: '%s', id: %d" % (self.type, str(self.name), str(self.value), self.id)) #+ "\n    "+str(self.t)

    def parse(self, data, header, hfile):
        global g_strings

        f = hfile
        pos = f.tell()

        t = struct.unpack(">iiii", data)
        self.type = t[1]
        atype = self.type
        self.handle = t[0]

        if atype==1: #int
            self.value = t[2]

        elif atype==2: #float
            t = struct.unpack(">iif4x", data)
            self.value = t[2]

        elif atype==3: #string
            t = struct.unpack(">iiii", data)
            self.offset = t[2]
            self.size = t[3]
            f.seek(header.stringtable_offset+self.offset)
            #self.value = unicode(f.read(self.size), 'utf-8')
            self.value = f.read(self.size)

        elif atype==6: #filename
            t = struct.unpack(">iiii", data)
            self.fileoffset = t[2]
            self.filesize = t[3]

        elif atype==7: #id
            t = struct.unpack(">iii4x", data)
            self.offset = t[2]
            f.seek(header.idtable_offset+self.offset)
            id_bin = f.read(4)
            (self.id,) = struct.unpack('>i', id_bin)
            self.value = ""
            a = f.read(1)
            while a!='\x00':
                self.value = self.value + str(a)
                a = f.read(1)

        f.seek(header.stringtable_offset+self.handle)
        self.name = ""
        a = f.read(1)
        while a!='\x00':
            self.name = self.name + str(a)
            a = f.read(1)
        f.seek(pos)

        self.t = t


class P3TExtractor:
    '''
    Basic format of a p3t theme file:

      - header (64)
      - tree
      - idtable
      - stringtable
      - intarraytable
      - floatarraytable
      - filetable

    '''
    def __init__(self, filename):
        self.themefile = filename
        self.header = P3THeader()

        try:
            p3tcompiler = zipimport.zipimporter('p3tcompiler.exe')
            self.cxml = p3tcompiler.load_module('cxml')
            self.header_size = self.cxml.header_bin_size
            self.header_fmt = ">"+self.cxml.header_bin_fmt
            self.element_size = self.cxml.element_bin_size
            self.element_fmt = ">"+self.cxml.element_bin_fmt
            self.elemap = {}
        except:
            self.header_size = 64
            self.header_fmt = ">4siiiiiiiiiiiii8x"
            self.element_size = 28
            self.element_fmt = ">iiiiiii"
            pass

        self.attr_fmt = ">iii4x"
        self.attr_size = struct.calcsize(self.attr_fmt)

    def parse(self):
        global g_strings

        h = self.header
        f = open(self.themefile, "rb")

        # parse header
        header_bin = f.read(self.header_size)
        (h.magic, h.version, h.tree_offset, h.tree_size,
         h.idtable_offset, h.idtable_size, h.stringtable_offset,
         h.stringtable_size, h.intarray_offset, h.intarray_size,
         h.floatarray_offset, h.floatarray_size, h.filetable_offset,
         h.filetable_size) = struct.unpack(self.header_fmt, header_bin)

        # load string table
        f.seek(h.stringtable_offset)
        stringtable_bin = f.read(h.stringtable_size)
        temp_strings = stringtable_bin.split('\x00')
        for a in temp_strings:
            g_strings.append(unicode(a, 'utf8'))

        # parse element tree
        f.seek(h.tree_offset)
        ele_size = self.element_size

        # read and parse the tree
        parentele = None
        while (f.tell()-h.tree_offset) < h.tree_size:
            element_bin = f.read(ele_size)
            ele = P3TElement(h)
            ele.parse(element_bin, self.element_fmt, f)
            
            # print ele

            # parse and add attributes to the element
            for y in range(0, ele.numattr):
                attr_bin = f.read(self.attr_size)
                attr = P3TAttribute()
                attr.parse(attr_bin, h, f)
                ele.add_attribute(attr)
                # print "    "+str(attr)

            # add to element map
            self.elemap[ele.offset] = ele

        # assign parents
        for k in self.elemap.keys():
            v = self.elemap[k]
            if k==0:
                continue
            if v.parentoffset==0:
                self.elemap[0].add_subelement(v)
            else:
                self.elemap[v.parentoffset].add_subelement(v)

        # restructure root element
        eleroot = self.elemap[0]
        self.elemap.clear()
        self.elemap[0] = eleroot

        # close file
        f.close()
        
    def dump_files(self, rootdir="extracted"):
        f = open(self.themefile, "rb")
        self.elemap[0].dump_files(rootdir, f)
        f.close()

    def dump_xml(self, rootdir="extracted"):
        document = xml.dom.minidom.Document()
        document.appendChild(self.elemap[0])
        f = open(rootdir+os.path.sep+os.path.basename(self.themefile)+".xml", "wb")
        f.write(document.toprettyxml())
        f.close()
        
def main():
    destination = "extracted"
    try:
        srcfile = sys.argv[1]
        try:
            destination = sys.argv[2]
        except:
            pass
    except:
        usage()
        return

    try:
        extr = P3TExtractor(srcfile)
        extr.parse()
        extr.dump_files(destination)
        extr.dump_xml(destination)
        print extr.header
    except Exception, e:
        usage()
        print str(e)
        traceback.print_exc()


def usage():
    print '''
P3T Unpacker %s
Copyright (c) 2007. Anoop Menon <codelogic at gmail dot com>

This program unpacks Playstation 3 Theme files (.p3t). By
default, it will extract the contents of the theme file
to the directory 'extracted' in the current directory.

It will also output an XML file (named <themename>.xml) which
can be used to recreate the P3T file using the theme compiler,
with the extracted images. The XML file will have the correct
filenames, so no renaming will need to be done.

This program is still in alpha stage and probably has dozens
of bugs. Feel free to fix them if you want.

If you have any suggestions, feel free to email me.

Usage:
   p3textractor <input theme file> [destination path]
''' % __version__

if __name__=="__main__":
    main()
